# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: donggu

from functools import wraps

import numpy as np
import itertools

# example:
# @vectorize
# # def f(a,b,c):
#   return a+b+c
#
# f([1,2], 0, [3,4]) == [[4,5],[5,6]]


def vectorize(f):
  @wraps(f)
  def wrapper(*args, **kwargs):
    items = []
    for (i, arg) in enumerate(args):
      if isinstance(arg, tuple) or isinstance(arg, list):
        items.append(tuple((i, j, a) for (j, a) in enumerate(arg)))

    for (k, value) in kwargs.items():
      if isinstance(value, tuple) or isinstance(value, list):
        items.append(tuple((k, i, v) for (i, v) in enumerate(value)))

    if items:
      # copy and modify this
      args = list(args)
      kwargs = kwargs.copy()

      outs = []
      for combination in itertools.product(*items):

        for i, config in enumerate(combination):
          if isinstance(config[0], int):
            args[config[0]] = config[2]
          else:
            kwargs[config[0]] = config[2]

        outs.append(f(*args, **kwargs))
        shape = [len(s) for s in items]

      outs = np.reshape(outs, shape).tolist()
      return outs
    else:
      return f(*args, **kwargs)

  return wrapper


__all__ = ['vectorize']
